# -*- coding: utf-8 -*-

# @Time : 2022/3/20 14:31

# @Author : Aweo
# @File : qt_test.py
import qdarkstyle
from PyQt5 import uic
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtWidgets import QApplication
import cv2
from PyQt5.QtCore import Qt
from nano_predict import Read
from threading import Thread

class Stats():
    def __init__(self):
        self.ui = uic.loadUi("ui_pro/my_ui.ui")
        self.ui.frame.setAlignment(Qt.AlignHCenter)  # 设置QLabel居中显示
        self.ui.BT1.clicked.connect(self.start_detect)
        self.ui.BT2.clicked.connect(self.close_detect)
        self.read = Read(label=['感染', '未感染'],
                         model_path='jetson_model/CNN5v2.pth',
                         model='CNN5')
        self.flag = None
        self.capture = cv2.VideoCapture(0)
        bg = QPixmap('ui_pro/img/438868d5da02722066c147c4f072a695.png')
        self.ui.bg_img.setPixmap(bg)
        self.ui.frame.setText('请打开摄像头')
        self.ui.frame.setAlignment(Qt.AlignCenter)

    def single_img_detect(self):
        img = cv2.imread('img_data/nag_img/5000.png', cv2.IMREAD_COLOR)
        res = self.read.Precdict(img)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        show_img = QImage(img.data, img.shape[0], img.shape[1], QImage.Format_RGB888)
        self.ui.result.setText(res)
        self.ui.frame.setPixmap(QPixmap.fromImage(show_img))

    def cap_detect(self):
        self.flag = True
        while self.flag:
            _, img = self.capture.read(cv2.IMREAD_COLOR)
            res = self.read.Precdict(img)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = cv2.resize(img, (420, 420))
            show_img = QImage(img.data, img.shape[0], img.shape[1], QImage.Format_RGB888)
            self.ui.result.setText(res)
            self.ui.frame.setPixmap(QPixmap.fromImage(show_img))

    def close_detect(self):
        self.end_detect()


    def end_detect(self):
        self.flag = False

    def start_detect(self):
        thread1 = Thread(target=self.cap_detect)
        thread1.start()


app = QApplication([])
stats = Stats()
# 显示ui
stats.ui.show()
app.exec_()
